from tqdm import tqdm

from deepLearning.model import MyModel
from dataset import *
from torch import nn
from sklearn.metrics import f1_score

model = MyModel()
model = model.to(device)

train_dataset = MyDataset('train')
valid_dataset = MyDataset('valid')

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
validation_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)

criteria = nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=lr)


# inputs是字典类型的，定义一个辅助函数帮助to(device)
def to_device(dict_tensors):
    result_tensors = {}
    for key, value in dict_tensors.items():
        result_tensors[key] = value.to(device)
    return result_tensors

# 定义一个验证方法，获取到验证集的精准率和loss
def validate():
    model.eval()
    total_loss = 0.
    total_correct = 0

    y_pred = []
    y_real = []
    for inputs, targets in validation_loader:
        inputs, targets = to_device(inputs), targets.to(device)
        outputs = model(inputs)
        loss = criteria(outputs.view(-1), targets.float())
        total_loss += float(loss)
        correct_num = (((outputs >= 0.5).float() * 1).flatten() == targets).sum()
        total_correct += correct_num

        outputs = (outputs > 0.5).float()
        y_pred.extend(outputs.tolist())
        y_real.extend(targets.tolist())

    f1 = f1_score(y_true=y_real, y_pred=y_pred, average='weighted')

    return total_correct / len(valid_dataset), total_loss / len(valid_dataset), f1


def train():
    # 首先将模型调成训练模式
    model.train()

    # 定义几个变量，帮助打印loss
    total_loss = 0.
    # 记录步数
    step = 0

    # 记录在验证集上最好的准确率
    best_f1 = 0

    # 开始训练
    for epoch in tqdm(range(epochs), desc='epoch training'):
        model.train()
        for i, (inputs, targets) in enumerate(train_loader):
            # 从batch中拿到训练数据
            inputs, targets = to_device(inputs), targets.to(device)
            # 传入模型进行前向传递
            outputs = model(inputs)
            # 计算损失
            loss = criteria(outputs.view(-1), targets.float())
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

            total_loss += float(loss)
            step += 1

            if step % log_per_step == 0:
                print("Epoch {}/{}, Step: {}/{}, total loss:{:.4f}".format(epoch + 1, epochs, i, len(train_loader),
                                                                           total_loss))
                total_loss = 0

            del inputs, targets

        # 一个epoch后，使用过验证集进行验证
        accuracy, validation_loss, cur_f1 = validate()
        print("Epoch {}, accuracy: {:.4f}, validation loss: {:.4f}, f1: {:.4f}".format(epoch + 1, accuracy, validation_loss, cur_f1))
        torch.save(model, model_dir / f"model_{epoch}.pt")

        # 保存最好的模型
        if cur_f1 > best_f1:
            torch.save(model, model_dir / f"model_best.pt")
            best_f1 = cur_f1

def test():
    # 加载最好的模型，然后进行测试集的预测
    model = torch.load(model_dir / f"model_best.pt")
    model = model.eval()
    test_dataset = MyDataset('test')
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)

    total_correct = 0
    y_pred = []
    y_real = []
    for inputs, targets in test_loader:
        inputs, targets = to_device(inputs), targets.to(device)
        outputs = model(inputs)

        outputs = (outputs > 0.5).int().flatten().tolist()
        y_pred.extend(outputs)
        y_real.extend(targets.tolist())

    f1 = f1_score(y_true=y_real, y_pred=y_pred, average='weighted')

    print(f'test result f1: {f1}')

if __name__ == '__main__':
    train()
    test()
